import platform
import subprocess
import argparse
import os
import re

def get_processes():
    """获取所有进程及其位置信息"""
    system = platform.system()
    processes = []
    
    if system == "Windows":
        try:
            # 使用 WMIC 获取带路径的进程信息
            output = subprocess.check_output(
                ["wmic", "process", "get", "ProcessId,ExecutablePath,CommandLine,Name", "/format:csv"],
                text=True,
                encoding="utf-8",
                errors="replace",
                creationflags=subprocess.CREATE_NO_WINDOW
            )
            
            # 解析 CSV 输出
            lines = output.strip().split("\n")[1:]  # 跳过标题行
            for line in lines:
                if not line.strip():
                    continue
                parts = line.split(',', 3)
                if len(parts) >= 4:
                    node, pid, path, name = parts[0], parts[1], parts[2], parts[3]
                    # 清理名称中的引号
                    name = name.replace('"', '').strip()
                    # 获取可执行文件名
                    exe_name = os.path.basename(path) if path else "Unknown"
                    # 获取目录位置
                    location = os.path.dirname(path) if path else "Unknown"
                    
                    processes.append({
                        "pid": pid.strip(),
                        "name": name,
                        "exe": exe_name,
                        "path": path.strip(),
                        "location": location,
                        "cmd": path  # Windows 上命令行为可执行文件路径
                    })
        except Exception as e:
            print(f"获取进程列表失败: {e}")
            return []
    
    elif system == "Linux":
        try:
            # 使用 ps 命令获取进程信息
            output = subprocess.check_output(
                ["ps", "-e", "-o", "pid,comm,args"],
                text=True,
                encoding="utf-8",
                errors="replace"
            )
            
            # 解析输出
            lines = output.strip().split("\n")[1:]  # 跳过标题行
            for line in lines:
                parts = line.strip().split(maxsplit=2)
                if len(parts) >= 3:
                    pid, comm, args = parts
                    
                    # 尝试获取实际路径
                    exe_path = ""
                    try:
                        exe_path = os.readlink(f"/proc/{pid}/exe")
                    except (FileNotFoundError, PermissionError, OSError):
                        exe_path = args.split()[0] if args else comm
                    
                    # 获取位置信息
                    location = os.path.dirname(exe_path) if exe_path else "Unknown"
                    exe_name = os.path.basename(exe_path) if exe_path else comm
                    
                    processes.append({
                        "pid": pid,
                        "name": comm,
                        "exe": exe_name,
                        "path": exe_path,
                        "location": location,
                        "cmd": args
                    })
        except Exception as e:
            print(f"获取进程列表失败: {e}")
            return []
    
    else:
        print(f"不支持的操作系统: {system}")
        return []
    
    return processes

def kill_process(pid):
    """杀死指定 PID 的进程"""
    system = platform.system()
    
    try:
        if system == "Windows":
            # 使用 taskkill 命令终止进程
            subprocess.run(
                ["taskkill", "/F", "/PID", str(pid)],
                check=True,
                creationflags=subprocess.CREATE_NO_WINDOW
            )
            print(f"✅ 进程 {pid} 已终止")
        
        elif system == "Linux":
            # 使用 kill 命令终止进程
            subprocess.run(["kill", "-9", str(pid)], check=True)
            print(f"✅ 进程 {pid} 已终止")
        
        else:
            print(f"不支持的操作系统: {system}")
    
    except subprocess.CalledProcessError as e:
        print(f"❌ 无法终止进程 {pid}: {e}")
    except Exception as e:
        print(f"❌ 发生错误: {e}")

def kill_by_name(name, exact_match=False):
    """根据进程名终止进程"""
    processes = get_processes()
    killed = 0
    
    for p in processes:
        if exact_match:
            match = p["name"].lower() == name.lower() or p["exe"].lower() == name.lower()
        else:
            match = name.lower() in p["name"].lower() or name.lower() in p["exe"].lower()
        
        if match:
            print(f"找到进程: PID={p['pid']}, 名称={p['name']}, 路径={p['path']}")
            confirm = input(f"⚠️ 确定要终止此进程吗? (y/n): ").lower()
            if confirm == "y":
                kill_process(p["pid"])
                killed += 1
    
    if killed == 0:
        print(f"❌ 未找到匹配进程: {name}")

def display_processes(processes, show_all=False):
    """显示进程列表"""
    if not processes:
        print("未找到任何进程")
        return
    
    system = platform.system()
    
    if system == "Windows":
        print("PID\t名称\t\t\t执行文件\t\t位置")
        print("-" * 100)
        for p in processes:
            name = p["name"][:20] + "..." if len(p["name"]) > 20 else p["name"]
            exe = p["exe"][:15] if p["exe"] else "N/A"
            location = p["location"][:40] + "..." if len(p["location"]) > 40 else p["location"]
            print(f"{p['pid']}\t{name}\t{exe}\t{location}")
    
    elif system == "Linux":
        print("PID\t名称\t\t执行文件\t\t位置")
        print("-" * 100)
        for p in processes:
            name = p["name"][:10] + "..." if len(p["name"]) > 10 else p["name"]
            exe = p["exe"][:15] if p["exe"] else "N/A"
            location = p["location"][:40] + "..." if len(p["location"]) > 40 else p["location"]
            print(f"{p['pid']}\t{name}\t{exe}\t{location}")
    
    print(f"\n总进程数: {len(processes)}")

def main():
    parser = argparse.ArgumentParser(description="高级进程管理器", formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("-l", "--list", action="store_true", help="列出所有进程")
    parser.add_argument("-k", "--kill", type=int, metavar="PID", help="终止指定PID的进程")
    parser.add_argument("-n", "--name", metavar="NAME", help="根据名称终止进程")
    parser.add_argument("-e", "--exact", action="store_true", help="与--name一起使用，精确匹配进程名")
    parser.add_argument("-s", "--search", metavar="KEYWORD", help="搜索包含关键字的进程")
    
    args = parser.parse_args()
    
    if not any(vars(args).values()):
        parser.print_help()
        return
    
    processes = []
    
    if args.list or args.search:
        processes = get_processes()
        
        if args.search:
            # 过滤包含搜索关键字的进程
            keyword = args.search.lower()
            filtered = [
                p for p in processes 
                if (keyword in p["name"].lower() or 
                    keyword in p["exe"].lower() or 
                    keyword in p["path"].lower() or 
                    keyword in p["location"].lower() or 
                    keyword in p["cmd"].lower())
            ]
            processes = filtered
        
        display_processes(processes)
    
    if args.kill:
        confirm = input(f"⚠️ 确定要终止进程 {args.kill} 吗? (y/n): ").lower()
        if confirm == "y":
            kill_process(args.kill)
        else:
            print("操作已取消")
    
    if args.name:
        kill_by_name(args.name, args.exact)

if __name__ == "__main__":
    #python process_manager.py --list
    #python process_manager.py --kill 1234
    # 模糊匹配（包含 "chrome" 的进程）
    #python process_manager.py --name chrome

    # 精确匹配（仅匹配 "chrome.exe"）
    #python process_manager.py --name chrome.exe --exact
    # 搜索所有包含 "python" 的进程
    #python process_manager.py --search python

    # 搜索特定路径的进程
    #python process_manager.py --search "/usr/bin"

    main()
